import numpy as np
from numpy.typing import NDArray
import copy
from typing import Dict

def get_manhattan_distance(l1, l2):
    return np.abs(l1[0] - l2[0]) + np.abs(l1[1] - l2[1])

def merge_close_lines(lines, eps=1e-1):
    new_lines = copy.deepcopy(lines)
    if len(new_lines) < 2:
        return new_lines
    i = 0
    while i < len(new_lines) - 1:
        j = i + 1
        while j < len(new_lines):
            l1, l2 = new_lines[i].points, new_lines[j].points
            if get_manhattan_distance(l1[0, :], l2[0, :]) < eps:
                l1 = l1[::-1, :]
                find = True
            elif get_manhattan_distance(l1[0, :], l2[-1, :]) < eps:
                l1 = l1[::-1, :]
                l2 = l2[::-1, :]
                find = True
            elif get_manhattan_distance(l1[-1, :], l2[0, :]) < eps:
                find = True
            elif get_manhattan_distance(l1[-1, :], l2[-1, :]) < eps:
                l2 = l2[::-1, :]
                find = True
            else:
                find = False
            if find:
                l1 = l1[:-1, :]
                l1 = np.concatenate([l1, l2], axis=0)
                new_lines[i].points = l1
                new_lines.pop(j)
                break
            else:
                j += 1
        if find:
            find = False
        else:
            i += 1
    return new_lines

def divide_lines(l1: NDArray, l2: NDArray, eps=1e-1):
    c1 = np.concatenate([l1[1:, :2].reshape(-1, 1, 2), l1[:-1, :2].reshape(-1, 1, 2)], axis=1)
    c2 = np.concatenate([l2[1:, :2].reshape(-1, 1, 2), l2[:-1, :2].reshape(-1, 1, 2)], axis=1)
    c1 = c1.reshape(-1, 1, 2, 2)
    c2 = c2.reshape(1, -1, 2, 2)
    diff1 = np.abs(c1 - c2).sum((2, 3)) <= eps
    diff2 = np.abs(c1 - c2[:, :, ::-1, :]).sum((2, 3)) <= eps
    mask = np.logical_or(diff1, diff2)
    coor1, coor2 = np.where(mask)
    lines_to_divide = [[], []]
    flag = False
    if len(coor1) > 0:
        coor1.sort()
        coor2.sort()
        coor1 = [-1] + coor1.tolist() + [l1.shape[0]]
        coor2 = [-1] + coor2.tolist() + [l2.shape[0]]
        for i in range(len(coor1)-1):
            lines_to_divide[0].append(l1[coor1[i]+1:coor1[i + 1]+1])
        for i in range(len(coor2)-1):
            lines_to_divide[1].append(l2[coor2[i]+1:coor2[i + 1]+1])
        flag = True
    return lines_to_divide, flag

def remove_duplicates(lines, eps=1e-1):
    new_lines = copy.deepcopy(lines)
    if len(new_lines) < 2:
        return new_lines
    i = 0
    find = False
    while i < len(new_lines) - 1:
        j = i + 1
        while j < len(new_lines):
            l1 = new_lines[i].points
            l2 = new_lines[j].points
            lines_to_divide, find = divide_lines(l1, l2, eps)
            if find:
                for new_l in lines_to_divide[0]:
                    new_item = copy.deepcopy(new_lines[i])
                    if new_l.shape[0] > 1:
                        new_item.points = new_l
                        new_lines.append(new_item)
                for new_l in lines_to_divide[1]:
                    new_item = copy.deepcopy(new_lines[j])
                    if new_l.shape[0] > 1:
                        new_item.points = new_l
                        new_lines.append(new_item)
                new_lines.pop(j)
                new_lines.pop(i)
                break
            else:
                j += 1
        if find:
            find = False
        else:
            i += 1
    return new_lines

def get_lane_segments_id_dict(lane_segments) -> Dict[int, int]:
    id_dict = {}
    for i, ls in enumerate(lane_segments):
        cur_lane = lane_segments[i]
        cur_id = cur_lane.id
        id_dict[cur_id] = i
    return id_dict
